import io
import logging
import time
import secrets
import uuid
from pathlib import Path
from typing import Union

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend


from common.safe_base64 import SafeBase64
from project.settings import COS_DIR, COS_BUCKET_FIRST_LEVER_DIRECTORY, QCLOUD_STORAGE_OPTION
from duchy.cert_pb2 import Cert
from common.cos import cos_client


class Theseus:
    DEFAULT_KEY_PREFIX = "fmZpZXJ5ZXllc35nb29kAA~~"
    DEFAULT_NONCE = "nVHPvy2K_n4lNtREHx10-g~~"

    def __init__(self, machine_code: str):
        self.default_key_prefix: bytes = SafeBase64.decode(self.DEFAULT_KEY_PREFIX)
        self.default_nonce: bytes = SafeBase64.decode(self.DEFAULT_NONCE)
        self.machine_code: str = machine_code
        self.mcb: bytes = SafeBase64.decode(machine_code)

    cert: Cert = None

    def decode_cert(self, cert_data: str) -> 'Cert':
        # TODO 异常捕获, 增强健壮性
        # 这里要拆开证书数据和加密tag
        cert_data, tag = cert_data.split('.')
        # 安全base64解码
        cert_byte: bytes = SafeBase64.decode(cert_data)
        aes_key: bytes = self.default_key_prefix[:16] + self.mcb[:16]
        nonce: bytes = self.mcb[16: 32]
        tag_byte: bytes = SafeBase64.decode(tag)
        # 初始化Cipher
        cipher = Cipher(algorithms.AES(aes_key), modes.GCM(nonce, tag=tag_byte), backend=default_backend())
        # 创建解密器
        decryptor = cipher.decryptor()
        # 解密数据, 同时获取解密之后的数据
        try:
            cert_dec: bytes = decryptor.update(cert_byte) + decryptor.finalize()
        except Exception as e:
            return None
        # TODO 异常捕获, 增强健壮性
        # 反序列化
        cert = Cert()
        cert.ParseFromString(cert_dec)  # 此处作反序列化, 是对cert本身做了反序列化
        self.cert = cert
        return self.cert

    # 加密证书
    def encode_cert(self, cert: Cert) -> str:
        # 这里应该是序列化 TODO 验证
        cert_byte: bytes = cert.SerializeToString()
        aes_key: bytes = self.default_key_prefix[:16] + self.mcb[:16]
        nonce: bytes = self.mcb[16: 32]
        # TODO 异常捕获, 增强健壮性
        # 初始化Cipher, cryptography默认标签长度为128位
        cipher = Cipher(algorithms.AES(aes_key), modes.GCM(nonce), backend=default_backend())
        # 创建加密器
        encryptor = cipher.encryptor()
        # 加密数据, 同时获取加密后的数据
        cert_enc: bytes = encryptor.update(cert_byte) + encryptor.finalize()
        # 安全base64编码
        cert_data: str = SafeBase64.encode(cert_enc)
        tag_data: str = SafeBase64.encode(encryptor.tag)
        return cert_data + "." + tag_data

    # 生成密钥
    def gen_key(self) -> bytes:
        # TODO python这边随机密钥和java采用不同的生成策略
        # 生成32字节的密钥
        return secrets.token_bytes(32)

    # ts文件加密
    def encode_ts(self, ts_data: bytes) -> bytes:

        key: bytes = self.cert.key

        aes_key: bytes = key[:16] + self.mcb[:16]
        nonce: bytes = self.mcb[16: 32]

        # TODO 异常捕获, 增强健壮性
        # 初始化Cipher
        cipher = Cipher(algorithms.AES(aes_key), modes.GCM(nonce), backend=default_backend())
        # 创建加密器
        encryptor = cipher.encryptor()
        # 加密数据, 并且获取加密后的数据
        ts_enc: bytes = encryptor.update(ts_data) + encryptor.finalize()
        final_enc_data = ts_enc + b'.' + encryptor.tag
        return final_enc_data


class TheseusTool:

    # 解密证书
    @staticmethod
    def decode_cert(args: list) -> 'Cert':
        # 去除两边空格
        machine_code: str = args[0].strip()
        cert_data: str = args[1].strip()
        t = Theseus(machine_code)
        cert: Cert = t.decode_cert(cert_data)
        return cert

    # 续期证书
    @staticmethod
    def renewal_cert(args: list, old_cert: Cert) -> str:
        machine_code: str = args[0].strip()
        expire_interval: int = args[1]
        # 获取当前时间戳(秒)
        ts_now = int(time.time())
        # 过期时间戳(秒)
        expire_time: int = ts_now + expire_interval
        t = Theseus(machine_code)
        cert: Cert = Cert()
        cert.create_time = old_cert.create_time
        cert.expire_time = expire_time
        cert.key = old_cert.key
        cert.extra = old_cert.extra
        # 加密证书
        cert_data: str = t.encode_cert(cert)
        return cert_data

    # 加密证书
    @staticmethod
    def encode_cert(args: list, original_key: Union[bytes, None]):
        machine_code: str = args[0].strip()
        expire_interval: int = int(args[1].strip())
        t = Theseus(machine_code)
        ts_now = int(time.time())
        # 到期时间戳
        expire_time: int = ts_now + expire_interval
        key: bytes = original_key if original_key else t.gen_key()
        cert = Cert()
        cert.create_time = ts_now
        cert.expire_time = expire_time
        cert.key = key
        cert.extra = "1|"
        cert_data: str = t.encode_cert(cert)
        return cert_data, cert

    # ts证书加密
    @staticmethod
    def encode_ts(args: list):
        machine_code: str = args[0].strip()
        cert_data: str = args[1].strip()
        ts_plain_file_path: str = args[2].strip()
        path = Path(ts_plain_file_path)
        # 读取ts文件的内容
        with open(path, 'rb') as f:
            ts_data: bytes = f.read()
        t = Theseus(machine_code)
        # 解密证书
        cert: Cert = t.decode_cert(cert_data)
        if not cert:
            return None, None
        # 加密ts
        ts_enc: bytes = t.encode_ts(ts_data)
        # 转成文件流
        file_stream = io.BytesIO(ts_enc)
        # 上传至cos
        file_path = COS_DIR + "/" + COS_BUCKET_FIRST_LEVER_DIRECTORY + "/" + str(uuid.uuid4()) + ".db"
        # TODO 补充异常捕获机制
        response = cos_client.upload_file_from_buffer(
            Bucket=QCLOUD_STORAGE_OPTION['Bucket'],
            Body=file_stream,
            Key=file_path,
            PartSize=1,
            MAXThread=10
        )
        print("file_path:", file_path)
        return file_path, ts_enc




